package com.liwy3.chart;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.theokanning.openai.completion.CompletionChoice;
import com.theokanning.openai.completion.CompletionRequest;
import com.theokanning.openai.completion.CompletionResult;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

/**
 * @author weiye.li
 * @date 2023/2/23
 */
public class SimpleChartDemo {

    private static final String API_ENDPOINT = "https://api.openai.com/v1/engines/text-davinci-003/completions";
    private static final String API_KEY = "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";

    public static void main(String[] args) throws IOException {

        StringBuilder msg = new StringBuilder();
        int count = 0;
        System.out.println("**********开始控制交互, 请输入问题**********");
        while (true) {
            count++;
            BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
            String prompt = reader.readLine();
            if (count == 10) {
                break;
            }
            msg.append(prompt).append("\n");
            String result = chart2Ai(msg.toString());
            System.out.println(result);
            System.out.println("==============AI回答完毕============");
        }
        System.out.println("**********控制台交互结束**********");
    }

    public static String chart2Ai(String prompt) {

        CompletionRequest completionRequest = CompletionRequest.builder()
                .prompt(prompt)
                .maxTokens(1024)
                //模型将承担风险的高低
                .temperature(0.9)
                .topP(1d)
                .frequencyPenalty(0.0)
                .presencePenalty(0.6)
                .build();
        String result = YQHttp.newRequest().setUrl(API_ENDPOINT)
                .addHeader("Authorization", "Bearer " + API_KEY)
                .addHeader("Content-Type", "application/json")
                .setPostBody(StrUtil.toUnderlineCase(JSONUtil.toJsonStr(completionRequest)))
                .setRequestConfig(YQHttp.RC_60)
                .json().getContent();
        CompletionResult completionResult = JSONUtil.toBean(result, CompletionResult.class);
        return completionResult.getChoices().stream().map(CompletionChoice::getText).findFirst().orElse("");
    }


}
